#define K_MAX_SHAPE_DIM 0

#include "kernel_operator.h"
using namespace AscendC;

// #define SCA
// #define PRINTF

class KernelEye {
   public:
    __aicore__ inline KernelEye() {}
    __aicore__ inline void Init(GM_ADDR y, GM_ADDR y_ref, uint16_t rows, uint16_t cols,
                                int32_t totalBatchNum, TPipe* pipeIn) {
        this->rows = rows;
        this->cols = cols;
        this->totalBatchNum = totalBatchNum;
        this->coreIndex = GetBlockIdx();
        this->coreNum = GetBlockNum();
        int32_t remainder = totalBatchNum % coreNum;
        this->batchNum = totalBatchNum / coreNum;
        this->batchBegin = coreIndex * batchNum + (coreIndex < remainder ? coreIndex : remainder);
        this->batchNum = batchNum + (coreIndex < remainder ? 1 : 0);
        this->setNum = rows < cols ? rows : cols;
        this->matrixSize = cols * rows;
#ifdef PRINTF
        printf("coreidx = %d ,[r,c] = %d %d , batch[begin, num] = %d %d\n",coreIndex,rows,cols,batchBegin,batchNum);
#endif
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y + rows * cols * batchBegin, batchNum * rows * cols);
#ifndef SCA
        this->pipe = pipeIn;
        pipe->InitBuffer(FillBuf, setNum * 32);
#endif
    }

    __aicore__ inline void Process() {
#ifdef SCA
        for (int32_t i = 0; i < batchNum; i++) {
            for (int32_t j = 0; j < setNum; j++) {
                yGm.SetValue(i * matrixSize + (cols + 1) * j, 1);
            }
        }
#else
        auto fillBuf = FillBuf.Get<DTYPE_Y>();
        uint32_t blockSize = 32 / sizeof(DTYPE_Y);
        if constexpr (std::is_same_v<DTYPE_Y, double>) {
        for (int32_t i = 0; i < setNum; i++) {
            fillBuf.SetValue(i * blockSize, (DTYPE_Y)1);
        } }
        else
        {
            Duplicate(fillBuf, (DTYPE_Y)1 , setNum * blockSize);
        }
        
        // DataCopyPadParams padParams{false, 0, 0, 0};
        uint32_t stride = cols * sizeof(DTYPE_Y);
        DataCopyExtParams copyParams{setNum, sizeof(DTYPE_Y), 0, stride,0};
        #pragma unroll 4
        for (int32_t i = 0; i < batchNum; i++) {
            DataCopyPad(yGm[i * matrixSize], fillBuf, copyParams);
        }

#endif
    }

   private:
    TPipe* pipe;
    GlobalTensor<DTYPE_Y> yGm;
    int32_t totalBatchNum, batchNum, batchBegin, matrixSize;
    uint32_t coreIndex, coreNum;
    uint16_t rows, cols, setNum;
    TBuf<QuePosition::VECCALC> FillBuf;
};

extern "C" __global__ __aicore__ void eye(GM_ADDR y, GM_ADDR y_ref, GM_ADDR workspace,
                                          GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    // KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
    KernelEye op;

    op.Init(y, y_ref, tiling_data.rows, tiling_data.cols, tiling_data.batchNum, &pipe);
    op.Process();
}